package cn.seecoder.ai.algorithm.ml.classification

import cn.seecoder.ai.algorithm.ml.BaseClassification
import cn.seecoder.ai.enums.TrainParamEnum
import cn.seecoder.ai.model.bo.TrainParamsBO
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.classification.GBTClassifier

import java.util

class GradientBoostedTreeClassification extends BaseClassification{


  override def buildMachineLearningStage(trainParams: TrainParamsBO): PipelineStage = {

    val gbt = new GBTClassifier()
      .setFeaturesCol("indexedFeatures")
      .setLabelCol("label")
      .setSeed(trainParams.seed)
      .setMaxDepth(trainParams.maxDepth)
      .setMaxBins(trainParams.maxBins)
      .setMinInstancesPerNode(trainParams.minInstancesPerNode)
      .setMinInfoGain(trainParams.minInfoGain)
      .setMaxIter(trainParams.maxIter)

    gbt
  }
}
